{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 32位AES\n",
    "\n",
    "在之前的AES教程（虽然是32位的平台）我们都是8位的操作模式，我们可以来看一些其他的ARM设备上的AES实现。\n",
    "\n",
    "你必须要有ARM版的板子才可以使用本教程，比如CWLite Arm或者UFO板+STM32F。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SCOPETYPE = 'OPENADC'\n",
    "PLATFORM = 'CWLITEARM'\n",
    "N = 1500\n",
    "CHECK_CORR = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 背景\n",
    "\n",
    "32位机器可以操作32位的字大小，所以每次只操作8位是非常浪费的，事实上，和AES的作者出版的[The Design of Rijndael](http://www.springer.com/gp/book/9783540425809)一书中所描述，我们可以生成几个表（T表）来加快AES速度。\n",
    "\n",
    "为了利用我们的32位机器，我们可以来考虑AES的工作流程，除了最后一轮外，其他每一轮都是：\n",
    "\n",
    "$\\text{a = Round Input}$\n",
    "\n",
    "$\\text{b = SubBytes(a)}$\n",
    "\n",
    "$\\text{c = ShiftRows(b)}$\n",
    "\n",
    "$\\text{d = MixColumns(c)}$\n",
    "\n",
    "$\\text{a' = AddRoundKey(d) = Round Output}$\n",
    "\n",
    "我们不管AddRoundKey操作，其他的操作分别是：\n",
    "\n",
    "$b_{i,j} = \\text{sbox}[a_{i,j}]$\n",
    "\n",
    "$\\left[ \\begin{array} { c } { c _ { 0 , j } } \\\\ { c _ { 1 , j } } \\\\ { c _ { 2 , j } } \\\\ { c _ { 3 , j } } \\end{array} \\right] = \\left[ \\begin{array} { l } { b _ { 0 , j + 0 } } \\\\ { b _ { 1 , j + 1 } } \\\\ { b _ { 2 , j + 2 } } \\\\ { b _ { 3 , j + 3 } } \\end{array} \\right]$\n",
    "\n",
    "$\\left[ \\begin{array} { l } { d _ { 0 , j } } \\\\ { d _ { 1 , j } } \\\\ { d _ { 2 , j } } \\\\ { d _ { 3 , j } } \\end{array} \\right] = \\left[ \\begin{array} { l l l l } { 02 } & { 03 } & { 01 } & { 01 } \\\\ { 01 } & { 02 } & { 03 } & { 01 } \\\\ { 01 } & { 01 } & { 02 } & { 03 } \\\\ { 03 } & { 01 } & { 01 } & { 02 } \\end{array} \\right] \\times \\left[ \\begin{array} { c } { c _ { 0 , j } } \\\\ { c _ { 1 , j } } \\\\ { c _ { 2 , j } } \\\\ { c _ { 3 , j } } \\end{array} \\right]$\n",
    "\n",
    "注意ShiftRows操作$b_{i, j+c}$是一个循环移位，MixColumns操作中的矩阵乘法是在$GF(2^8)$域上的xtime运算。\n",
    "\n",
    "我们可以把这些操作合并到一行中，我们可以将四个不同的4子节向量合并到$d$中：\n",
    "\n",
    "$\\left[ \\begin{array} { l } { d _ { 0 , j } } \\\\ { d _ { 1 , j } } \\\\ { d _ { 2 , j } } \\\\ { d _ { 3 , j } } \\end{array} \\right] = \\left[ \\begin{array} { l } { 02 } \\\\ { 01 } \\\\ { 01 } \\\\ { 03 } \\end{array} \\right] \\operatorname { sbox } \\left[ a _ { 0 , j + 0 } \\right] \\oplus \\left[ \\begin{array} { l } { 03 } \\\\ { 02 } \\\\ { 01 } \\\\ { 01 } \\end{array} \\right] \\operatorname { sbox } \\left[ a _ { 1 , j + 1 } \\right] \\oplus \\left[ \\begin{array} { c } { 01 } \\\\ { 03 } \\\\ { 02 } \\\\ { 01 } \\end{array} \\right] \\operatorname { sbox } \\left[ a _ { 2 , j + 2 } \\right] \\oplus \\left[ \\begin{array} { c } { 01 } \\\\ { 01 } \\\\ { 03 } \\\\ { 02 } \\end{array} \\right] \\operatorname { sbox } \\left[ a _ { 3 , j + 3 } \\right]$\n",
    "\n",
    "现在，对于这四个组件的每一个，我们都可以将每个可能的8位输入的输出制成表格\n",
    "\n",
    "$T _ { 0 } [ a ] = \\left[ \\begin{array} { l l } { 02 \\times \\operatorname { sbox } [ a ] } \\\\ { 01 \\times \\operatorname { sbox } [ a ] } \\\\ { 01 \\times \\operatorname { sbox } [ a ] } \\\\ { 03 \\times \\operatorname { sbox } [ a ] } \\end{array} \\right]$\n",
    "\n",
    "$T _ { 1 } [ a ] = \\left[ \\begin{array} { l } { 03 \\times \\operatorname { sbox } [ a ] } \\\\ { 02 \\times \\operatorname { sbox } [ a ] } \\\\ { 01 \\times \\operatorname { sbox } [ a ] } \\\\ { 01 \\times \\operatorname { sbox } [ a ] } \\end{array} \\right]$\n",
    "\n",
    "$T _ { 2 } [ a ] = \\left[ \\begin{array} { l l } { 01 \\times \\operatorname { sbox } [ a ] } \\\\ { 03 \\times \\operatorname { sbox } [ a ] } \\\\ { 02 \\times \\operatorname { sbox } [ a ] } \\\\ { 01 \\times \\operatorname { sbox } [ a ] } \\end{array} \\right]$\n",
    "\n",
    "$T _ { 3 } [ a ] = \\left[ \\begin{array} { l l } { 01 \\times \\operatorname { sbox } [ a ] } \\\\ { 01 \\times \\operatorname { sbox } [ a ] } \\\\ { 03 \\times \\operatorname { sbox } [ a ] } \\\\ { 02 \\times \\operatorname { sbox } [ a ] } \\end{array} \\right]$\n",
    "\n",
    "这些表格拥有$2^8$个不同的32位条目，所以这些表将总共占用4kB，最后，我们可以快速的计算出AES的一轮结果\n",
    "\n",
    "$\\left[ \\begin{array} { l } { d _ { 0 , j } } \\\\ { d _ { 1 , j } } \\\\ { d _ { 2 , j } } \\\\ { d _ { 3 , j } } \\end{array} \\right] = T _ { 0 } \\left[ a _ { 0 } , j + 0 \\right] \\oplus T _ { 1 } \\left[ a _ { 1 } , j + 1 \\right] \\oplus T _ { 2 } \\left[ a _ { 2 } , j + 2 \\right] \\oplus T _ { 3 } \\left[ a _ { 3 } , j + 3 \\right]$\n",
    "\n",
    "加上最后的AddRoundKey操作，现在我们一个回合需要16次查表操作和16次32的XOR操作，这种方式比传统的8位实现要高效很多，同时还可以进行一些折中，比如表仅存储8位位移，所以也可能只需要存储1kB的查询表，但是这样会牺牲一些旋转操作。\n",
    "\n",
    "注意从侧信道的角度来说，T表对于AES并没有多大的影响，这些SubBytes的输出依然隐藏在T表中，而其他的操作都是线性的，所以我们依然有可能使用8位的攻击方式对32位的AES发起攻击。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 硬件\n",
    "\n",
    "硬件和之前的AES例子类似，不过这次我们需要用到MBEDTLS进行构建（之前用的是TINYAES128C）："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "cd ../hardware/victims/firmware/\n",
    "mkdir -p simpleserial-aes-lab2 && cp -r simpleserial-aes/* $_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "CRYPTO_TARGET = \"MBEDTLS\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash -s \"$PLATFORM\" \"$CRYPTO_TARGET\"\n",
    "cd ../hardware/victims/firmware/simpleserial-aes-lab2\n",
    "make PLATFORM=$1 CRYPTO_TARGET=$2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 运行攻击\n",
    "\n",
    "攻击和分析和普通的AES攻击没有任何不同，因为其实我们并不关心T表的实现，我们只需要捕获更多的轨迹（500个），但我们的轨迹也没必要那么长（1500样本即可）"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 捕获轨迹"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run \"Helper_Scripts/Setup_Generic.ipynb\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fw_path = \"../hardware/victims/firmware/simpleserial-aes-lab2/simpleserial-aes-{}.hex\".format(PLATFORM)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cw.program_target(scope, prog, fw_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "project = cw.create_project(\"projects/AES32.cwp\", overwrite = True)\n",
    "\n",
    "#Capture Traces\n",
    "from tqdm import tnrange\n",
    "import numpy as np\n",
    "import time\n",
    "\n",
    "ktp = cw.ktp.Basic()\n",
    "scope.adc.samples = 1500\n",
    "\n",
    "for i in tnrange(N, desc='Capturing traces'):\n",
    "    key, text = ktp.next()  # manual creation of a key, text pair can be substituted here\n",
    "    trace = cw.capture_trace(scope, target, text, key)\n",
    "    if trace is None:\n",
    "        continue\n",
    "    project.traces.append(trace)\n",
    "project.save()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# cleanup the connection to the target and scope\n",
    "scope.dis()\n",
    "target.dis()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 分析"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import chipwhisperer.analyzer as cwa\n",
    "leak_model = cwa.leakage_models.sbox_output\n",
    "attack = cwa.cpa(project, leak_model)\n",
    "\n",
    "cb = cwa.get_jupyter_callback(attack)   \n",
    "attack_results = attack.run(cb, 100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "此时，你会看到表格顶部都是红色数字。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 绘制\n",
    "\n",
    "这次我们跳过大部分情节，但我们还是要看看输出/时间的关系："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bokeh.plotting import figure, show\n",
    "from bokeh.io import output_notebook\n",
    "plot_data = cwa.analyzer_plots(attack_results)\n",
    "\n",
    "output_notebook()\n",
    "rets = []\n",
    "for i in range(0, 16):\n",
    "    rets.append(plot_data.output_vs_time(i))\n",
    "\n",
    "p = figure()\n",
    "for ret in rets:\n",
    "    p.line(ret[0], ret[2], line_color='green')\n",
    "    p.line(ret[0], ret[3], line_color='green')\n",
    "    \n",
    "for ret in rets:\n",
    "    p.line(ret[0], ret[1], line_color='red')\n",
    "\n",
    "show(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "放大，你应该看到生成的图比8位的AES实现更加混乱。\n",
    "\n",
    "![](https://wiki.newae.com/images/e/e3/32bit_AES_outvstime.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "key = np.array(project.keys[0])\n",
    "recv_key = [kguess[0][0] for kguess in attack_results.find_maximums()]\n",
    "assert np.all((key == recv_key)), \"Failed to recover encryption key\\nGot: {}\\nExpected: {}\".format(recv_key, key)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert (attack_results.pge == [0]*16), \"PGE for some bytes not zero: {}\".format(attack_results.pge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# MBED Often has a few bytes with low (0.2-0.4) correlation, causing it to fail this test\n",
    "if CHECK_CORR:\n",
    "    max_corrs = [kguess[0][2] for kguess in attack_results.find_maximums()]\n",
    "    assert (np.all([corr > 0.6 for corr in max_corrs])), \"Low correlation in attack (corr <= 0.6): {}\".format(max_corrs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
